package rbtree;

public class RBTree {
    static class RBTreeNode {
        public RBTreeNode left;
        public RBTreeNode right;
        public RBTreeNode parent;
        public int val;
        public COLOR color;
        public RBTreeNode(int val) {
            this.val = val;
            //新增结点，默认必须是红色
            this.color = COLOR.RED;
        }
    }
    public RBTreeNode root;//根节点

    public boolean insert(int val) {
        RBTreeNode node = new RBTreeNode(val);
        if(root == null) {
            root = node;
            root.color = COLOR.BLACK;//注意第一次插入根节点一定是黑色
            return true;
        }
        RBTreeNode parent = null;
        RBTreeNode cur = root;
        while(cur != null) {
            if(cur.val > val) {
                //向左寻找
                parent = cur;
                cur = cur.left;
            } else if(cur.val == val) {
                //相等就说明插入失败
                return false;
            } else {
                //向右寻找
                parent = cur;
                cur = cur.right;
            }
        }
        //cur == null
        if(parent.val > val) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        node.parent = parent;
        cur = node;

        //红黑树需要调整颜色
        while (parent != null && parent.color == COLOR.RED) {//parent为红色，就是两个红色结点连在一起了
            RBTreeNode grandFather = parent.parent;//这个引用一定不为空，因为此时parent一定是红色！
            if(parent == grandFather.left) {
                RBTreeNode uncle = grandFather.right;
                if(uncle != null && uncle.color == COLOR.RED) {
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    //继续向上调整
                    cur = grandFather;
                    parent = cur.parent;
                } else {
                    //uncle不存在 或者 uncle是黑色的
                    //情况三：(这里只需要把情况三处理成情况二)
                    if(cur == parent.right) {
                        rotateLeft(parent);
                        //这里只需要交换一下cur和parent
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    }//走完这里，情况三就变成了情况二

                    //情况二：
                    //这里只需要先右旋，然后修改颜色即可
                    rotateRight(parent);
                    parent.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                }
            } else {
                //parent == grandFather.right
                RBTreeNode uncle = grandFather.left;
                if(uncle != null && uncle.color == COLOR.RED) {
                    parent.color = COLOR.BLACK;
                    uncle.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                    //继续向上调整
                    cur = grandFather;
                    parent = cur.parent;
                } else {
                    //uncle不存在 或者 uncle是黑色的
                    //情况三：(这里只需要把情况三处理成情况二)
                    if(cur == parent.left) {
                        rotateRight(parent);
                        //这里只需要交换一下cur和parent
                        RBTreeNode tmp = parent;
                        parent = cur;
                        cur = tmp;
                    }//走完这里，情况三就变成了情况二

                    //情况二：
                    //这里只需要先右旋，然后修改颜色即可
                    rotateLeft(parent);
                    parent.color = COLOR.BLACK;
                    grandFather.color = COLOR.RED;
                }
            }
        }
        root.color = COLOR.BLACK;
        return true;
    }

    //左单旋
    private void rotateLeft(RBTreeNode parent) {
        RBTreeNode subR = parent.right;
        RBTreeNode subRL = subR.left;
        parent.right = subRL;

        subR.left = parent;

        if(subRL != null) {
            subRL.parent = parent;
        }
        //必须先记录parent的父亲
        RBTreeNode pParent = parent.parent;
        parent.parent = subR;
        //检查parent是否为根节点
        if(parent == root) {
            root = subR;
            root.parent = null;
        } else {
            //不是根节点需要判断parent是左子树还是右子树
            if(pParent.left == parent) {
                pParent.left = subR;
            } else {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
    }

    //右单旋
    private void rotateRight(RBTreeNode parent) {
        RBTreeNode subL = parent.left;
        RBTreeNode subLR = subL.right;
        parent.left = subLR;
        subL.right = parent;
        if(subLR != null) {
            subLR.parent = parent;
        }
        //必须先记录parent的父亲
        RBTreeNode pParent = parent.parent;
        parent.parent = subL;
        //检查parent是否是根节点
        if(parent == root) {
            root = subL;
            root = null;
        } else {
            //不是根节点，判断parent是左子树还是右子树
            if(pParent.left == parent) {
                pParent.left = subL;
            } else {
                pParent.right = subL;
            }
            subL.parent = pParent;
        }
    }

    //验证红黑树(就是要满足红黑树的性质)
    public boolean isRBTree(RBTreeNode root) {
        if(root == null) {
            return true;
        }
        if (root.color != COLOR.BLACK) {
            System.out.println("不是红黑树，违反了性质：根节点是黑色的");
        }
        int BlackNum = 0;
        RBTreeNode cur = root;
        while(root != null) {
            if(root.color == COLOR.BLACK) {
                BlackNum++;
            }
            root = root.left;
        }
        // 检查是否出现俩个连续的红色结点 && 检查每条路径黑色结点数是否相同
        return checkRedColor(root) && checkBlackNum(root, 0, BlackNum);
    }
    /**
     *
     * @param root
     * @param pathBlackNum 递归时每条路径上的黑色结点个数
     * @param BlackNum 以及计算好的黑色结点个数
     * @return
     */
    private boolean checkBlackNum(RBTreeNode root, int pathBlackNum, int BlackNum) {
        if(root == null) {
            return true;
        }
        if(root.color == COLOR.BLACK) {
            pathBlackNum++;
        }
        if(root.left == null && root.right == null) {
            if(pathBlackNum != BlackNum) {
                System.out.println("不是红黑树，违反了性质：每条路径上的黑色结点个数相同");
                return false;
            }
        }
        return checkBlackNum(root.left, pathBlackNum, BlackNum) &&
                checkBlackNum(root.left, pathBlackNum, BlackNum);

    }
    //检查是否出现俩个连续的红色结点
    private boolean checkRedColor(RBTreeNode root) {
        if(root == null) {
            return true;
        }
        if(root.color == COLOR.RED) {
            RBTreeNode parent = root.parent;
            if(parent.color == COLOR.RED) {
                System.out.println("不是红黑树，违反了性质：不可以出现两个连续的红色结点");
                return false;
            }
        }
        return checkRedColor(root.left) && checkRedColor(root.right);
    }
    //检查中序遍历是否有序
    public void dfs(RBTreeNode root) {
        if(root == null) {
            return;
        }
        dfs(root.left);
        System.out.print(root.val + " ");
        dfs(root.right);
    }
}